
#include <torch/script.h>
#include <opencv2/highgui/highgui.hpp>
#include <opencv4/opencv2/imgcodecs/imgcodecs.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/opencv.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <iostream>
#include <memory>
#include <string>
void classify(cv::Mat image)
{
    cv::Mat img = image.clone();
    auto device = torch::Device(torch::kCPU);
    torch::jit::script::Module module = torch::jit::load("/home/lenovo/_test/nn/pokemon/model.pt");
    std::cout << "Load model successful!" << std::endl;
    module.to(device);
    module.eval();
    std::vector<int64_t> sizes = {1, img.rows, img.cols, 3};
    torch::Tensor img_tensor = torch::from_blob(img.data, torch::IntList(sizes), torch::kByte);
    // img_tensor = torch::rand(0,3,100,100);
    img_tensor = img_tensor.permute({0, 3, 1, 2}).toType(torch::kFloat32);
    img_tensor = img_tensor.div(255); //将上述数字到0——1
    // std::cout << img_tensor[0] << std::endl;

    img_tensor[0][0] = img_tensor[0][0].sub_(0.485).div_(0.229);
    img_tensor[0][1] = img_tensor[0][1].sub_(0.456).div_(0.224);
    img_tensor[0][2] = img_tensor[0][2].sub_(0.406).div_(0.225);

    std::cout << img_tensor[0] << std::endl;
    torch::Tensor output = module.forward({img_tensor}).toTensor();
    output = torch::softmax(output, 1);
    // auto max_result = output.max(1, true);
    // auto max_index = std::get<1>(max_result).item<float>();
    // std::cout << output.size << std::endl;
    // std::cout << max_result << std::endl;
    // std::cout << max_index << std::endl;
    int prevalue = torch::argmax(output).item<int>();
    std::cout << prevalue << std::endl;
}
int main()
{
    cv::Mat img = cv::imread("/home/lenovo/_test/nn/pokemon/test/3/q22.png");
    cv::resize(img, img, cv::Size(100, 100), cv::COLOR_BGR2RGB);
    std::cout << img.size() << std::endl;
    classify(img);
    return 0;
}
